#include "header/AES.h"
AES::AES(int keyLen)
{
    this->Nb = 4;
    switch (keyLen)
    {
        case 128:
            this->Nk = 4;
            this->Nr = 10;
            break;
        case 192:
            this->Nk = 6;
            this->Nr = 12;
            break;
        case 256:
            this->Nk = 8;
            this->Nr = 14;
            break;
        default:
            throw "Incorrect key length";
    }

    blockBytesLen = 4 * this->Nb * sizeof(unsigned char);
}



unsigned char *AES::EncryptCBC(unsigned char in[], unsigned int inLen, unsigned  char key[], unsigned char * iv, unsigned int &outLen)
{
    outLen = GetPaddingLength(inLen);
    unsigned char *alignIn  = PaddingNulls(in, inLen, outLen);
    //unsigned char *out = new unsigned char[outLen];
    unsigned char *out = new unsigned char[outLen+1];//这里后面必须多加一个\0结束
    out[outLen] = '\0';
    unsigned char *block = new unsigned char[blockBytesLen];
    memcpy(block, iv, blockBytesLen);
    for (unsigned int i = 0; i < outLen; i+= blockBytesLen)
    {
        XorBlocks(block, alignIn + i, block, blockBytesLen);
        EncryptBlock(block, out + i, key);
        memcpy(block, out + i, blockBytesLen);
    }

    delete[] block;
    delete[] alignIn;

    return out;
}

//解码其实输入串都是对齐了的，是blockSize(16)的整数倍，inLen就是outLen
unsigned char *AES::DecryptCBC(unsigned char in[], unsigned int inLen, unsigned  char key[], unsigned char * iv)
{
    //unsigned char *out = new unsigned char[inLen];
    unsigned char *out = new unsigned char[inLen];
    unsigned char *block = new unsigned char[blockBytesLen];
    memcpy(block, iv, blockBytesLen);
    for (unsigned int i = 0; i < inLen; i+= blockBytesLen)
    {
        DecryptBlock(in + i, out + i, key);
        XorBlocks(block, out + i, out + i, blockBytesLen);
        memcpy(block, in + i, blockBytesLen);
    }
    //移除补齐位
    removePaddingNulls(out, inLen);
    delete[] block;

    return out;
}

static const unsigned  char HEX[16]={0x10,0x01,0x02,0x03,0x04,0x05,0x06,0x07,0x08,0x09,0x0a,0x0b,0x0c,0x0d,0x0e,0x0f};
int *AES::findPaddingIndex(unsigned char *str, size_t length) {
    static int result[] = {-1, -1}, i, k;
    for (i = 0; i < length; ++i) {
        char c = str[length - i];
        if ('\0' != c) {
            result[0] = i;
            for (k = 0; k < blockBytesLen; ++k) {
                if (HEX[k] == c) {
                    if (0 == k) {
                        k = blockBytesLen;
                    }
                    result[1] = k;
                    return result;
                }
            }
            return result;
        }
    }
}
/**
 * 移除填充
 */
void AES::removePaddingNulls(unsigned char out[],const size_t inputLength){
    int *result = findPaddingIndex(out, inputLength - 1);
    int offSetIndex = result[0];
    int lastChar = result[1];
    //检查是不是padding的字符,然后去掉
    const size_t noZeroIndex = inputLength - offSetIndex;
    if (lastChar >= 0 && offSetIndex >= 0) {
        int success = 1;
        for (int i = 0; i < lastChar; ++i) {
            size_t index = noZeroIndex - lastChar + i;
            if (!HEX[lastChar] == out[index]) {
                success = 0;
            }
        }
        if (1 == success) {
            out[noZeroIndex - lastChar] = '\0';
            memset(out + noZeroIndex - lastChar + 1, 0, lastChar - 1);
        }
    } else {
        out[noZeroIndex] = '\0';
    }
}

/**
 * <1>已对齐，填充一个长度为blockSize(16)且每个字节均为blockSize的数据。
  <2>未对齐，需要补充的字节个数为n，则填充一个长度为n且每个字节均为n的数据。
 * @param in
 * @param inLen
 * @param alignLen
 * @return
 */
unsigned char * AES::PaddingNulls(unsigned char in[], unsigned int inLen, unsigned int alignLen)
{
    /*unsigned char *alignIn = new unsigned char[alignLen];
    memcpy(alignIn, in, inLen);
    memset(alignIn + inLen, 0x00, alignLen - inLen);//这里不能直接赋0
    return alignIn;*/
    int k=inLen%16;
    int nPidding_size = 16 - k;
    unsigned char *alignIn = new unsigned char[alignLen];
    memcpy(alignIn, in, inLen);
    memset(alignIn + inLen, nPidding_size, alignLen - inLen);
    return alignIn;
}
/**
 * * <1>已对齐，填充一个长度为blockSize(16)且每个字节均为blockSize的数据。
  <2>未对齐，需要补充的字节个数为n，则填充一个长度为n且每个字节均为n的数据。
 */
unsigned int AES::GetPaddingLength(unsigned int len)
{
    unsigned int lengthWithPadding =  (len / blockBytesLen);
    //if (len % blockBytesLen) {//不要if了，无论如何都要补充
        lengthWithPadding++;
    //}
    lengthWithPadding *=  blockBytesLen;
    return lengthWithPadding;
}

void AES::EncryptBlock(unsigned char in[], unsigned char out[], unsigned  char key[])
{
    unsigned char *w = new unsigned char[4 * Nb * (Nr + 1)];
    KeyExpansion(key, w);
    unsigned char **state = new unsigned char *[4];
    state[0] = new unsigned  char[4 * Nb];
    int i, j, round;
    for (i = 0; i < 4; i++)
    {
        state[i] = state[0] + Nb * i;
    }


    for (i = 0; i < 4; i++)
    {
        for (j = 0; j < Nb; j++)
        {
            state[i][j] = in[i + 4 * j];
        }
    }

    AddRoundKey(state, w);

    for (round = 1; round <= Nr - 1; round++)
    {
        SubBytes(state);
        ShiftRows(state);
        MixColumns(state);
        AddRoundKey(state, w + round * 4 * Nb);
    }

    SubBytes(state);
    ShiftRows(state);
    AddRoundKey(state, w + Nr * 4 * Nb);

    for (i = 0; i < 4; i++)
    {
        for (j = 0; j < Nb; j++)
        {
            out[i + 4 * j] = state[i][j];
        }
    }

    delete[] state[0];
    delete[] state;
    delete[] w;
}

void AES::DecryptBlock(unsigned char in[], unsigned char out[], unsigned  char key[])
{
    unsigned char *w = new unsigned char[4 * Nb * (Nr + 1)];
    KeyExpansion(key, w);
    unsigned char **state = new unsigned char *[4];
    state[0] = new unsigned  char[4 * Nb];
    int i, j, round;
    for (i = 0; i < 4; i++)
    {
        state[i] = state[0] + Nb * i;
    }


    for (i = 0; i < 4; i++)
    {
        for (j = 0; j < Nb; j++) {
            state[i][j] = in[i + 4 * j];
        }
    }

    AddRoundKey(state, w + Nr * 4 * Nb);

    for (round = Nr - 1; round >= 1; round--)
    {
        InvSubBytes(state);
        InvShiftRows(state);
        AddRoundKey(state, w + round * 4 * Nb);
        InvMixColumns(state);
    }

    InvSubBytes(state);
    InvShiftRows(state);
    AddRoundKey(state, w);

    for (i = 0; i < 4; i++)
    {
        for (j = 0; j < Nb; j++) {
            out[i + 4 * j] = state[i][j];
        }
    }

    delete[] state[0];
    delete[] state;
    delete[] w;
}


void AES::SubBytes(unsigned char **state)
{
    int i, j;
    unsigned char t;
    for (i = 0; i < 4; i++)
    {
        for (j = 0; j < Nb; j++)
        {
            t = state[i][j];
            state[i][j] = sbox[t / 16][t % 16];
        }
    }

}

void AES::ShiftRow(unsigned char **state, int i, int n)    // shift row i on n positions
{
    unsigned char t;
    int k, j;
    for (k = 0; k < n; k++)
    {
        t = state[i][0];
        for (j = 0; j < Nb - 1; j++)
        {
            state[i][j] = state[i][j + 1];
        }
        state[i][Nb - 1] = t;
    }
}

void AES::ShiftRows(unsigned char **state)
{
    ShiftRow(state, 1, 1);
    ShiftRow(state, 2, 2);
    ShiftRow(state, 3, 3);
}

unsigned char AES::xtime(unsigned char b)    // multiply on x
{
    unsigned char mask = 0x80, m = 0x1b;
    unsigned char high_bit = b & mask;
    b = b << 1;
    if (high_bit) {    // mod m(x)
        b = b ^ m;
    }
    return b;
}

unsigned char AES::mul_bytes(unsigned char a, unsigned char b)
{
    unsigned char c = 0, mask = 1, bit, d;
    int i, j;
    for (i = 0; i < 8; i++)
    {
        bit = b & mask;
        if (bit)
        {
            d = a;
            for (j = 0; j < i; j++)
            {    // multiply on x^i
                d = xtime(d);
            }
            c = c ^ d;    // xor to result
        }
        b = b >> 1;
    }
    return c;
}

void AES::MixColumns(unsigned char **state)
{
    unsigned char s[4], s1[4];
    int i, j;

    for (j = 0; j < Nb; j++)
    {
        for (i = 0; i < 4; i++)
        {
            s[i] = state[i][j];
        }

        s1[0] = mul_bytes(0x02, s[0]) ^ mul_bytes(0x03, s[1]) ^ s[2] ^ s[3];
        s1[1] = s[0] ^ mul_bytes(0x02, s[1]) ^ mul_bytes(0x03, s[2]) ^ s[3];
        s1[2] = s[0] ^ s[1] ^ mul_bytes(0x02, s[2]) ^ mul_bytes(0x03, s[3]);
        s1[3] = mul_bytes(0x03, s[0]) ^ s[1] ^ s[2] ^ mul_bytes(0x02, s[3]);
        for (i = 0; i < 4; i++)
        {
            state[i][j] = s1[i];
        }

    }

}

void AES::AddRoundKey(unsigned char **state, unsigned char *key)
{
    int i, j;
    for (i = 0; i < 4; i++)
    {
        for (j = 0; j < Nb; j++)
        {
            state[i][j] = state[i][j] ^ key[i + 4 * j];
        }
    }
}

void AES::SubWord(unsigned char *a)
{
    int i;
    for (i = 0; i < 4; i++)
    {
        a[i] = sbox[a[i] / 16][a[i] % 16];
    }
}

void AES::RotWord(unsigned char *a)
{
    unsigned char c = a[0];
    a[0] = a[1];
    a[1] = a[2];
    a[2] = a[3];
    a[3] = c;
}

void AES::XorWords(unsigned char *a, unsigned char *b, unsigned char *c)
{
    int i;
    for (i = 0; i < 4; i++)
    {
        c[i] = a[i] ^ b[i];
    }
}

void AES::Rcon(unsigned char * a, int n)
{
    int i;
    unsigned char c = 1;
    for (i = 0; i < n - 1; i++)
    {
        c = xtime(c);
    }

    a[0] = c;
    a[1] = a[2] = a[3] = 0;
}

void AES::KeyExpansion(unsigned char key[], unsigned char w[])
{
    unsigned char *temp = new unsigned char[4];
    unsigned char *rcon = new unsigned char[4];

    int i = 0;
    while (i < 4 * Nk)
    {
        w[i] = key[i];
        i++;
    }

    i = 4 * Nk;
    while (i < 4 * Nb * (Nr + 1))
    {
        temp[0] = w[i - 4 + 0];
        temp[1] = w[i - 4 + 1];
        temp[2] = w[i - 4 + 2];
        temp[3] = w[i - 4 + 3];

        if (i / 4 % Nk == 0)
        {
            RotWord(temp);
            SubWord(temp);
            Rcon(rcon, i / (Nk * 4));
            XorWords(temp, rcon, temp);
        }
        else if (Nk > 6 && i / 4 % Nk == 4)
        {
            SubWord(temp);
        }

        w[i + 0] = w[i - 4 * Nk] ^ temp[0];
        w[i + 1] = w[i + 1 - 4 * Nk] ^ temp[1];
        w[i + 2] = w[i + 2 - 4 * Nk] ^ temp[2];
        w[i + 3] = w[i + 3 - 4 * Nk] ^ temp[3];
        i += 4;
    }

    delete []rcon;
    delete []temp;

}


void AES::InvSubBytes(unsigned char **state)
{
    int i, j;
    unsigned char t;
    for (i = 0; i < 4; i++)
    {
        for (j = 0; j < Nb; j++)
        {
            t = state[i][j];
            state[i][j] = inv_sbox[t / 16][t % 16];
        }
    }
}

void AES::InvMixColumns(unsigned char **state)
{
    unsigned char s[4], s1[4];
    int i, j;

    for (j = 0; j < Nb; j++)
    {
        for (i = 0; i < 4; i++)
        {
            s[i] = state[i][j];
        }
        s1[0] = mul_bytes(0x0e, s[0]) ^ mul_bytes(0x0b, s[1]) ^ mul_bytes(0x0d, s[2]) ^ mul_bytes(0x09, s[3]);
        s1[1] = mul_bytes(0x09, s[0]) ^ mul_bytes(0x0e, s[1]) ^ mul_bytes(0x0b, s[2]) ^ mul_bytes(0x0d, s[3]);
        s1[2] = mul_bytes(0x0d, s[0]) ^ mul_bytes(0x09, s[1]) ^ mul_bytes(0x0e, s[2]) ^ mul_bytes(0x0b, s[3]);
        s1[3] = mul_bytes(0x0b, s[0]) ^ mul_bytes(0x0d, s[1]) ^ mul_bytes(0x09, s[2]) ^ mul_bytes(0x0e, s[3]);

        for (i = 0; i < 4; i++)
        {
            state[i][j] = s1[i];
        }
    }
}

void AES::InvShiftRows(unsigned char **state)
{
    ShiftRow(state, 1, Nb - 1);
    ShiftRow(state, 2, Nb - 2);
    ShiftRow(state, 3, Nb - 3);
}

void AES::XorBlocks(unsigned char *a, unsigned char * b, unsigned char *c, unsigned int len)
{
    for (unsigned int i = 0; i < len; i++)
    {
        c[i] = a[i] ^ b[i];
    }
}

